{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc69fc7d",
   "metadata": {},
   "source": [
    "# `relay.qnn.requantize`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bed6bee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import te\n",
    "import numpy as np\n",
    "from tvm import relay\n",
    "from tvm.contrib import graph_executor\n",
    "\n",
    "# 定义舍入模式：向上舍入和四舍五入\n",
    "roundings = [\"UPWARD\", \"TONEAREST\"]\n",
    "# 定义计算数据类型\n",
    "compute_dtypes = [\"float32\", \"float64\", \"int64\"]\n",
    "# 定义输出量化数据类型\n",
    "out_dtypes = [\"int8\", \"int16\"]\n",
    "\n",
    "\n",
    "def verify(mod, goldens, target=\"llvm\"):\n",
    "    \"\"\"\n",
    "    验证QNN重量化模型的输出是否符合预期\n",
    "\n",
    "    参数:\n",
    "    mod: 要验证的Relay模块\n",
    "    goldens: 包含输入数据和预期输出的元组\n",
    "    target: 目标设备，默认为\"llvm\"\n",
    "    \"\"\"\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        # 构建模型\n",
    "        graph, lib, params = relay.build(mod, target, params=None)\n",
    "        golden_data, golden_output = goldens\n",
    "        # 创建执行器\n",
    "        rt_mod = graph_executor.create(graph, lib, device=tvm.cpu(0))\n",
    "        # 设置输入数据\n",
    "        rt_mod.set_input(\"input_data\", golden_data)\n",
    "        rt_mod.set_input(**params)\n",
    "        # 运行模型\n",
    "        rt_mod.run()\n",
    "        # 获取输出\n",
    "        res = rt_mod.get_output(0).numpy()\n",
    "        # 验证结果\n",
    "        np.testing.assert_equal(res, golden_output)\n",
    "\n",
    "\n",
    "def get_mod(\n",
    "    data_shape,\n",
    "    data_dtype,\n",
    "    out_dtype,\n",
    "    input_scale,\n",
    "    output_scale,\n",
    "    input_zero_point=0,\n",
    "    output_zero_point=0,\n",
    "    rounding=\"None\",\n",
    "    compute_dtype=\"None\",\n",
    "    axis=0,\n",
    "):\n",
    "    \"\"\"\n",
    "    创建QNN重量化模型的辅助函数\n",
    "\n",
    "    参数:\n",
    "    data_shape: 输入数据形状\n",
    "    data_dtype: 输入数据类型\n",
    "    out_dtype: 输出数据类型\n",
    "    input_scale: 输入缩放比例\n",
    "    output_scale: 输出缩放比例\n",
    "    input_zero_point: 输入零点偏移，默认为0\n",
    "    output_zero_point: 输出零点偏移，默认为0\n",
    "    rounding: 舍入模式，默认为\"None\"\n",
    "    compute_dtype: 计算数据类型，默认为\"None\"\n",
    "    axis: 通道轴，默认为0\n",
    "\n",
    "    返回:\n",
    "    创建的Relay模块\n",
    "    \"\"\"\n",
    "    # 创建输入变量\n",
    "    input_data = relay.var(\"input_data\", shape=data_shape, dtype=data_dtype)\n",
    "    # 处理输入缩放比例\n",
    "    if isinstance(input_scale, float):\n",
    "        input_scale_expr = relay.const(input_scale, \"float32\")\n",
    "    else:\n",
    "        input_scale_expr = relay.const(np.array(input_scale).astype(\"float32\"))\n",
    "\n",
    "    # 处理输入零点偏移\n",
    "    if isinstance(input_zero_point, float):\n",
    "        input_zero_point_expr = relay.const(input_zero_point, \"int32\")\n",
    "    else:\n",
    "        input_zero_point_expr = relay.const(np.array(input_zero_point).astype(\"int32\"))\n",
    "\n",
    "    # 创建重量化操作\n",
    "    mod = relay.qnn.requantize(\n",
    "        input_data,\n",
    "        input_scale=input_scale_expr,\n",
    "        input_zero_point=input_zero_point_expr,\n",
    "        output_scale=relay.const(output_scale, \"float32\"),\n",
    "        output_zero_point=relay.const(output_zero_point, \"int32\"),\n",
    "        axis=axis,\n",
    "        rounding=rounding,\n",
    "        compute_dtype=compute_dtype,\n",
    "        out_dtype=out_dtype,\n",
    "    )\n",
    "\n",
    "    # 构建函数和模块\n",
    "    mod = relay.Function(relay.analysis.free_vars(mod), mod)\n",
    "    mod = tvm.IRModule.from_expr(mod)\n",
    "    return mod\n",
    "\n",
    "\n",
    "def test_same_scale():\n",
    "    \"\"\"\n",
    "    测试输入和输出缩放比例相同的情况\n",
    "\n",
    "    当输入和输出缩放比例相同时，重量化操作应该是一个恒等变换\n",
    "    \"\"\"\n",
    "    # 生成测试数据：范围在[-100, 100)的整数\n",
    "    golden_data = np.arange(-100, 100, 1).astype(\"int32\")\n",
    "    # 预期输出与输入相同\n",
    "    golden_output = golden_data\n",
    "    # 测试所有计算数据类型、舍入模式和输出数据类型的组合\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型\n",
    "                mod = get_mod(\n",
    "                    data_shape=(200,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=0.5,\n",
    "                    output_scale=0.5,  # 相同的缩放比例\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "                # 确保没有右移操作（因为缩放比例相同）\n",
    "                assert \"right_shift\" not in mod.astext()\n",
    "                # 验证结果\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_scalar_same_scale():\n",
    "    \"\"\"\n",
    "    测试标量输入且缩放比例相同的情况\n",
    "\n",
    "    验证对标量数据进行重量化时的正确性\n",
    "    \"\"\"\n",
    "    # 生成标量测试数据\n",
    "    golden_data = np.array(-10).astype(\"int32\")\n",
    "    # 预期输出与输入相同\n",
    "    golden_output = golden_data\n",
    "    # 测试所有计算数据类型、舍入模式和输出数据类型的组合\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（标量输入）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(),  # 标量形状\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=0.5,\n",
    "                    output_scale=0.5,  # 相同的缩放比例\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "                # 确保没有右移操作\n",
    "                assert \"right_shift\" not in mod.astext()\n",
    "                # 验证结果\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_downscale():\n",
    "    \"\"\"\n",
    "    测试输出缩放比例大于输入缩放比例（缩小）的情况\n",
    "\n",
    "    验证当输出缩放比例大于输入缩放比例时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（缩放比例 1->16）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=16,  # 输出缩放比例更大（缩小）\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                # 8对应0.5，结果为1\n",
    "                golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "                golden_output = np.repeat([0, 1, 2], [8, 16, 8])\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                # -8对应-0.5。对于UPWARD舍入，结果为0\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "                if rounding == \"UPWARD\":\n",
    "                    golden_output = np.repeat([0, -1, -2], [9, 16, 7])\n",
    "                else:\n",
    "                    golden_output = np.repeat([0, -1, -2], [8, 16, 8])\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试不同的缩放比例（1->4）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=4,\n",
    "                    rounding=rounding,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                # 2I对应0.5，结果为1\n",
    "                golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "                golden_output = np.repeat([0, 1, 2, 3, 4, 5, 6, 7, 8], [2, 4, 4, 4, 4, 4, 4, 4, 2])\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                # -8对应-0.5。对于UPWARD舍入，结果为0\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "                if rounding == \"UPWARD\":\n",
    "                    golden_output = np.repeat(\n",
    "                        [0, -1, -2, -3, -4, -5, -6, -7, -8], [3, 4, 4, 4, 4, 4, 4, 4, 1]\n",
    "                    )\n",
    "                else:\n",
    "                    golden_output = np.repeat(\n",
    "                        [0, -1, -2, -3, -4, -5, -6, -7, -8], [2, 4, 4, 4, 4, 4, 4, 4, 2]\n",
    "                    )\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "            # 测试uint8输出数据类型\n",
    "            mod = get_mod(\n",
    "                data_shape=(32,),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"uint8\",\n",
    "                input_scale=1,\n",
    "                output_scale=16,\n",
    "                rounding=rounding,\n",
    "            )\n",
    "\n",
    "            # 测试正值\n",
    "            golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "            golden_output = np.repeat([0, 1, 2], [8, 16, 8])\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "            # 测试uint8输入和输出数据类型\n",
    "            mod = get_mod(\n",
    "                data_shape=(32,),\n",
    "                data_dtype=\"uint8\",\n",
    "                out_dtype=\"uint8\",\n",
    "                input_scale=1,\n",
    "                output_scale=16,\n",
    "                rounding=rounding,\n",
    "            )\n",
    "\n",
    "            # 测试正值\n",
    "            golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "            golden_output = np.repeat([0, 1, 2], [8, 16, 8])\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_upscale():\n",
    "    \"\"\"\n",
    "    测试输出缩放比例小于输入缩放比例（放大）的情况\n",
    "\n",
    "    验证当输出缩放比例小于输入缩放比例时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（缩放比例 2->1）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=2,\n",
    "                    output_scale=1,  # 输出缩放比例更小（放大）\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "                # 由于输出缩放比例是输入的一半，输出值应该是输入值的2倍\n",
    "                golden_output = np.multiply(2, golden_data)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "                golden_output = np.multiply(2, golden_data)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_non_power_of_two():\n",
    "    \"\"\"\n",
    "    测试非2的幂次缩放比例的情况\n",
    "\n",
    "    验证当缩放比例不是2的幂次时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（缩放比例 1->3）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=3,  # 非2的幂次缩放比例\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                # 输入数据乘以3，以便重量化后得到原始范围\n",
    "                golden_data = np.multiply(np.arange(0, 32, 1).astype(\"int32\"), 3)\n",
    "                golden_output = np.arange(0, 32, 1)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                golden_data = np.multiply(np.arange(0, -32, -1).astype(\"int32\"), 3)\n",
    "                golden_output = np.arange(0, -32, -1)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试不同的缩放比例（3->1）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=3,\n",
    "                    output_scale=1,\n",
    "                    rounding=rounding,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "                # 由于输出缩放比例是输入的1/3，输出值应该是输入值的3倍\n",
    "                golden_output = np.multiply(golden_data, 3)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "                golden_output = np.multiply(golden_data, 3)\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_saturation_int8():\n",
    "    \"\"\"\n",
    "    测试int8数据类型的饱和情况\n",
    "\n",
    "    验证当重量化结果超出int8类型范围时的饱和行为\n",
    "    \"\"\"\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            # 创建重量化模型（输出类型为int8）\n",
    "            mod = get_mod(\n",
    "                data_shape=(16,),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=0.5,\n",
    "                output_scale=0.5,\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            # 测试正值饱和\n",
    "            golden_data = np.arange(0, 16, 1).astype(\"int32\")\n",
    "            golden_data = np.add(120, golden_data)  # 从120开始，超出int8的最大值127\n",
    "            output = np.array(\n",
    "                [120, 121, 122, 123, 124, 125, 126, 127, 127, 127, 127, 127, 127, 127, 127, 127]\n",
    "            )\n",
    "            golden_output = output\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "            # 测试负值饱和\n",
    "            golden_data = np.arange(0, -16, -1).astype(\"int32\")\n",
    "            golden_data = np.add(-120, golden_data)  # 从-120开始，超出int8的最小值-128\n",
    "            output = np.array(\n",
    "                [\n",
    "                    -120,\n",
    "                    -121,\n",
    "                    -122,\n",
    "                    -123,\n",
    "                    -124,\n",
    "                    -125,\n",
    "                    -126,\n",
    "                    -127,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                    -128,\n",
    "                ]\n",
    "            )\n",
    "            golden_output = output\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_saturation_int16():\n",
    "    \"\"\"\n",
    "    测试int16数据类型的饱和情况\n",
    "\n",
    "    验证当重量化结果超出int16类型范围时的饱和行为\n",
    "    \"\"\"\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            # 创建重量化模型（输出类型为int16）\n",
    "            mod = get_mod(\n",
    "                data_shape=(16,),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int16\",\n",
    "                input_scale=0.5,\n",
    "                output_scale=0.5,\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            # 测试正值饱和\n",
    "            golden_data = np.arange(0, 16, 1).astype(\"int32\")\n",
    "            golden_data = np.add(32760, golden_data)  # 从32760开始，超出int16的最大值32767\n",
    "            output = np.array(\n",
    "                [\n",
    "                    32760,\n",
    "                    32761,\n",
    "                    32762,\n",
    "                    32763,\n",
    "                    32764,\n",
    "                    32765,\n",
    "                    32766,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                    32767,\n",
    "                ]\n",
    "            )\n",
    "            golden_output = output\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "            # 测试负值饱和\n",
    "            golden_data = np.arange(0, -16, -1).astype(\"int32\")\n",
    "            golden_data = np.add(-32760, golden_data)  # 从-32760开始，超出int16的最小值-32768\n",
    "            output = np.array(\n",
    "                [\n",
    "                    -32760,\n",
    "                    -32761,\n",
    "                    -32762,\n",
    "                    -32763,\n",
    "                    -32764,\n",
    "                    -32765,\n",
    "                    -32766,\n",
    "                    -32767,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                    -32768,\n",
    "                ]\n",
    "            )\n",
    "            golden_output = output\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_zero_point():\n",
    "    \"\"\"\n",
    "    测试零点偏移对重量化的影响\n",
    "\n",
    "    分别测试输出零点偏移和输入零点偏移的情况\n",
    "    \"\"\"\n",
    "    # 测试输出零点偏移\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            # 创建重量化模型（输出零点偏移为1）\n",
    "            mod = get_mod(\n",
    "                data_shape=(32,),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=1,\n",
    "                output_scale=16,\n",
    "                output_zero_point=1,  # 输出零点偏移\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "\n",
    "            # 测试正值\n",
    "            golden_data = np.arange(0, 32, 1).astype(\"int32\")\n",
    "            golden_output = np.repeat([0, 1, 2], [8, 16, 8])\n",
    "            golden_output = np.add(1, golden_output)  # 加上输出零点偏移\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "            # 测试负值\n",
    "            golden_data = np.arange(-32, -64, -1).astype(\"int32\")\n",
    "            if rounding == \"UPWARD\":\n",
    "                golden_output = np.repeat([-2, -3, -4], [9, 16, 7])\n",
    "            else:\n",
    "                golden_output = np.repeat([-2, -3, -4], [8, 16, 8])\n",
    "            golden_output = np.add(1, golden_output)  # 加上输出零点偏移\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "    # 测试输入零点偏移\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（输入零点偏移为16）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=16,\n",
    "                    input_zero_point=16,  # 输入零点偏移\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "\n",
    "                # 测试正值\n",
    "                golden_data = np.arange(32, 64, 1).astype(\"int32\")\n",
    "                golden_output = np.repeat([2, 3, 4], [8, 16, 8])\n",
    "                golden_output = np.subtract(golden_output, 1)  # 减去输入零点偏移的影响\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "                # 测试负值\n",
    "                golden_data = np.arange(-32, -64, -1).astype(\"int32\")\n",
    "                if rounding == \"UPWARD\":\n",
    "                    golden_output = np.repeat([-2, -3, -4], [9, 16, 7])\n",
    "                else:\n",
    "                    golden_output = np.repeat([-2, -3, -4], [8, 16, 8])\n",
    "                golden_output = np.subtract(golden_output, 1)  # 减去输入零点偏移的影响\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_per_channel_same_scale():\n",
    "    \"\"\"\n",
    "    测试通道级相同缩放比例的情况\n",
    "\n",
    "    验证当每个通道使用相同的缩放比例时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    # 测试2D数据\n",
    "    golden_data = np.arange(-5, 5, 1).astype(\"int32\").reshape((5, 2))\n",
    "    golden_output = golden_data  # 预期输出与输入相同\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 创建重量化模型（每个通道使用相同的缩放比例）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(5, 2),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=[0.5, 0.5],  # 通道级缩放比例\n",
    "                    output_scale=0.5,\n",
    "                    axis=1,  # 通道轴\n",
    "                    rounding=rounding,\n",
    "                    compute_dtype=compute_dtype,\n",
    "                )\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "    # 测试3D数据（改变通道轴）\n",
    "    golden_data = np.arange(-10, 10, 1).astype(\"int32\").reshape((2, 2, 5))\n",
    "    golden_output = golden_data\n",
    "\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            # 创建重量化模型（不同的输入形状和通道轴）\n",
    "            mod = get_mod(\n",
    "                data_shape=(2, 2, 5),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=[0.5, 0.5],  # 通道级缩放比例\n",
    "                output_scale=0.5,\n",
    "                axis=1,  # 通道轴\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_per_channel_different_scale():\n",
    "    \"\"\"\n",
    "    测试通道级不同缩放比例的情况\n",
    "\n",
    "    验证当每个通道使用不同的缩放比例时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    # 测试2D数据（不同通道缩放比例）\n",
    "    golden_data = np.arange(-5, 5, 1).astype(\"int32\").reshape((5, 2))\n",
    "    golden_output = np.array([-5, -2, -3, -1, -1, 0, 1, 1, 3, 2]).reshape((5, 2))\n",
    "\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            # 创建重量化模型（每个通道使用不同的缩放比例）\n",
    "            mod = get_mod(\n",
    "                data_shape=(5, 2),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=[0.5, 0.25],  # 不同的通道缩放比例\n",
    "                output_scale=0.5,\n",
    "                axis=1,\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "    # 测试3D数据（改变通道轴）\n",
    "    golden_data = np.arange(-20, 20, 2).astype(\"int32\").reshape((2, 2, 5))\n",
    "    golden_output = np.array(\n",
    "        [-20, -18, -16, -14, -12, -5, -4, -3, -2, -1, 0, 2, 4, 6, 8, 5, 6, 7, 8, 9]\n",
    "    ).reshape((2, 2, 5))\n",
    "\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            mod = get_mod(\n",
    "                data_shape=(2, 2, 5),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=[0.5, 0.25],\n",
    "                output_scale=0.5,\n",
    "                axis=1,\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "    # 测试输入缩放比例大于输出缩放比例的情况\n",
    "    golden_data = np.arange(-5, 5, 1).astype(\"int32\").reshape((5, 2))\n",
    "    golden_output = np.array([-10, -2, -6, -1, -2, 0, 2, 1, 6, 2]).reshape((5, 2))\n",
    "\n",
    "    for compute_dtype in compute_dtypes:\n",
    "        for rounding in roundings:\n",
    "            mod = get_mod(\n",
    "                data_shape=(5, 2),\n",
    "                data_dtype=\"int32\",\n",
    "                out_dtype=\"int8\",\n",
    "                input_scale=[1.0, 0.25],  # 输入缩放比例大于输出\n",
    "                output_scale=0.5,\n",
    "                axis=1,\n",
    "                rounding=rounding,\n",
    "                compute_dtype=compute_dtype,\n",
    "            )\n",
    "            verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_default_cfg_and_no_args():\n",
    "    \"\"\"\n",
    "    测试默认配置且无参数的情况\n",
    "\n",
    "    验证使用默认配置且不指定额外参数时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for qnn_out_dtype in out_dtypes:\n",
    "        # 创建重量化模型（使用默认配置）\n",
    "        mod = get_mod(\n",
    "            data_shape=(32,),\n",
    "            data_dtype=\"int32\",\n",
    "            out_dtype=qnn_out_dtype,\n",
    "            input_scale=1,\n",
    "            output_scale=16,\n",
    "            # 不指定舍入模式等参数，使用默认值\n",
    "        )\n",
    "        golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "        golden_output = np.repeat([0, -1, -2], [9, 16, 7])\n",
    "        verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_non_default_cfg_and_no_args():\n",
    "    \"\"\"\n",
    "    测试非默认配置且无参数的情况\n",
    "\n",
    "    验证使用非默认全局配置且不指定额外参数时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for rounding_cfg in roundings:\n",
    "        for qnn_out_dtype in out_dtypes:\n",
    "            # 设置非默认全局配置\n",
    "            with relay.qnn.requantize_config(rounding=rounding_cfg):\n",
    "                # 创建重量化模型（不指定额外参数）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=16,\n",
    "                )\n",
    "\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "\n",
    "                # 根据全局配置的舍入模式设置预期输出\n",
    "                if rounding_cfg == \"UPWARD\":\n",
    "                    golden_output = np.repeat([0, -1, -2], [9, 16, 7])\n",
    "                else:\n",
    "                    golden_output = np.repeat([0, -1, -2], [8, 16, 8])\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_default_cfg_and_args():\n",
    "    \"\"\"\n",
    "    测试默认配置且有参数的情况\n",
    "\n",
    "    验证使用默认全局配置但指定局部参数时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for rounding in roundings:\n",
    "        for qnn_out_dtype in out_dtypes:\n",
    "            # 设置默认全局配置\n",
    "            with relay.qnn.requantize_config(rounding=\"UPWARD\"):\n",
    "                # 创建重量化模型（指定局部参数，覆盖全局配置）\n",
    "                mod = get_mod(\n",
    "                    data_shape=(32,),\n",
    "                    data_dtype=\"int32\",\n",
    "                    out_dtype=qnn_out_dtype,\n",
    "                    input_scale=1,\n",
    "                    output_scale=16,\n",
    "                    rounding=rounding,  # 局部参数覆盖全局配置\n",
    "                )\n",
    "\n",
    "                golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "\n",
    "                # 根据局部参数的舍入模式设置预期输出\n",
    "                if rounding == \"UPWARD\":\n",
    "                    golden_output = np.repeat([0, -1, -2], [9, 16, 7])\n",
    "                else:\n",
    "                    golden_output = np.repeat([0, -1, -2], [8, 16, 8])\n",
    "                verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "def test_non_default_cfg_and_args():\n",
    "    \"\"\"\n",
    "    测试非默认配置且有参数的情况\n",
    "\n",
    "    验证使用非默认全局配置且指定局部参数时，重量化操作的正确性\n",
    "    \"\"\"\n",
    "    for rounding_arg in roundings:\n",
    "        for rounding_cfg in roundings:\n",
    "            for qnn_out_dtype in out_dtypes:\n",
    "                # 设置非默认全局配置\n",
    "                with relay.qnn.requantize_config(rounding=rounding_cfg):\n",
    "                    # 创建重量化模型（指定局部参数，覆盖全局配置）\n",
    "                    mod = get_mod(\n",
    "                        data_shape=(32,),\n",
    "                        data_dtype=\"int32\",\n",
    "                        out_dtype=qnn_out_dtype,\n",
    "                        input_scale=1,\n",
    "                        output_scale=16,\n",
    "                        rounding=rounding_arg,  # 局部参数覆盖全局配置\n",
    "                    )\n",
    "\n",
    "                    golden_data = np.arange(0, -32, -1).astype(\"int32\")\n",
    "\n",
    "                    # 根据局部参数的舍入模式设置预期输出\n",
    "                    if rounding_arg == \"UPWARD\":\n",
    "                        golden_output = np.repeat([0, -1, -2], [9, 16, 7])\n",
    "                    else:\n",
    "                        golden_output = np.repeat([0, -1, -2], [8, 16, 8])\n",
    "                    verify(mod, (golden_data, golden_output))\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # 执行所有测试函数\n",
    "    test_same_scale()  # 测试相同缩放比例\n",
    "    test_scalar_same_scale()  # 测试标量相同缩放比例\n",
    "    test_downscale()  # 测试缩小缩放比例\n",
    "    test_upscale()  # 测试放大缩放比例\n",
    "    test_non_power_of_two()  # 测试非2的幂次缩放比例\n",
    "    test_saturation_int8()  # 测试int8饱和情况\n",
    "    test_saturation_int16()  # 测试int16饱和情况\n",
    "    test_zero_point()  # 测试零点偏移\n",
    "    test_per_channel_same_scale()  # 测试通道级相同缩放比例\n",
    "    test_per_channel_different_scale()  # 测试通道级不同缩放比例\n",
    "    test_default_cfg_and_no_args()  # 测试默认配置且无参数\n",
    "    test_non_default_cfg_and_no_args()  # 测试非默认配置且无参数\n",
    "    test_default_cfg_and_args()  # 测试默认配置且有参数\n",
    "    test_non_default_cfg_and_args()  # 测试非默认配置且有参数\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
